#include "worker.h"
#include "log.h"
#include "define.h"

#include <iostream>
#include <string>
#include <arpa/inet.h>

static void _Run_Loop(void *arg)
{
    Worker *worker = (Worker *)arg;

    worker->RunLoop();

    return;
}

static void _RecvSessionConnection(uv_stream_t *pServer, int status)
{
    LOG_DEBUG("_RecvSessionConnection %d %p", status, pServer);
    if (status < 0 || pServer == NULL)
    {
        return;
    }

    Worker *pWorker = (Worker *)uv_handle_get_data((uv_handle_t *)pServer);
    if (pWorker == NULL)
    {
        uv_close((uv_handle_t *)pServer, NULL);
        return;
    }
    pWorker->OnRecvSessionConnection(pServer);

    return;
}

static void _RecvSpiderConnection(uv_stream_t *pServer, int status)
{
    if (status < 0 || pServer == NULL)
    {
        return;
    }
    Worker *pWorker = (Worker *)uv_handle_get_data((uv_handle_t *)pServer);
    if (pWorker == NULL)
    {
        uv_close((uv_handle_t *)pServer, NULL);
        return;
    }
    pWorker->OnRecvSpiderConnection(pServer);
    return;
}

static void _OnAllocBuff(uv_handle_t *handle, size_t suggested_size, uv_buf_t *buf)
{
    SpiderBackConn *pSpiderBackConn = (SpiderBackConn *)uv_handle_get_data(handle);

    buf->base = pSpiderBackConn->m_cmdRecvBuf;
    buf->len = 4;
    return;
}

static void _OnReadData(uv_stream_t *stream, ssize_t nread, const uv_buf_t *buf)
{
    LOG_DEBUG("SpiderBackConn _OnReadData");
    SpiderBackConn *pSpiderBackConn = (SpiderBackConn *)uv_handle_get_data((uv_handle_t *)stream);
    if (pSpiderBackConn == nullptr)
    {
        uv_close((uv_handle_t *)stream, NULL);
        return;
    }
    if (nread == 4)
    {
        pSpiderBackConn->DoConn();
    }
    else
    {
        pSpiderBackConn->Close();
    }
    return;
}

static void _OnSpiderBackConnTimeout(uv_timer_t *handle) {
    SpiderBackConn *pSpiderBackConn = (SpiderBackConn *)uv_handle_get_data((uv_handle_t *)handle);
    if (pSpiderBackConn) {
        pSpiderBackConn->Close();
    }
}


void _SpiderBackConnClose(uv_handle_t* handle){
    SpiderBackConn *pSession = (SpiderBackConn *)uv_handle_get_data(handle);
    if (pSession) {
        delete pSession;
    }
}

bool Worker::Init(int index, int workerType, int listenPortStart, int listenPortEnd, int spiderPort, int commandTimeout, int idleTimeout)
{
    m_index = index;
    m_sessionId = m_index << 24 | 1;
    LOG_DEBUG("Worker: Init m_sessionId %u", m_sessionId);
    m_workerType = workerType;
    m_listenPortStart = listenPortStart;
    m_listenPortEnd = listenPortEnd;
    m_listenNum = m_listenPortEnd - m_listenPortStart + 1;
    m_spiderPort = spiderPort;
    m_commandTimeout = commandTimeout;
    m_idleTimeout = idleTimeout;
    if (m_workerType == WORKER_TYPE_SPIDER && m_spiderPort == 0)
    {
        LOG_ERROR("Worker %d Init parm is error", m_index);
        return false;
    }
    LOG_INFO("Worker %d Init listenPort:%d-%d spiderPort:%d", m_index, m_listenPortStart, m_listenPortEnd, m_spiderPort);
    return true;
}

bool Worker::AddIdentification(const std::string username, const std::string password)
{
    if (username.empty() || password.empty())
    {
        return true;
    }
    m_identification[username] = password;
    return true;
}

bool Worker::AddWhiteList(const std::string ipMask)
{
    uint32_t ip = 0;
    uint32_t mask = 32;
    std::string ipStr;
    size_t pos = ipMask.find("/");
    if (pos != std::string::npos)
    {
        ipStr = ipMask.substr(0, pos);
        std::string maskStr = ipMask.substr(pos + 1);
        mask = std::stoi(maskStr);
    }
    else
    {
        ipStr = ipMask;
    }
    struct in_addr addr;
    if (inet_pton(AF_INET, ipStr.c_str(), &addr) == 1)
    {
        ip = addr.s_addr;
    }
    else
    {
        LOG_ERROR("Add WhiteList IP[%s] is error", ipStr.c_str());
        return false;
    }
    LOG_DEBUG("Add WhiteList ipMask[%s] [%u] IP [%08x] MASK[%08x]", ipMask.c_str(), mask, ip, g_MaskTrans[mask]);
    for (auto &maskIps : m_whiteList)
    {
        if (maskIps.mask == g_MaskTrans[mask])
        {
            maskIps.IPTbl.insert((ip & maskIps.mask));
            return true;
        }
    }
    MaskIpSet entry;
    entry.mask = g_MaskTrans[mask];
    entry.IPTbl.insert(ip & entry.mask);
    m_whiteList.emplace_back(entry);
    return true;
}

bool Worker::DoAuth(const std::string username, const std::string password, uint32_t ipaddr)
{
    bool ret = true;
    for (const auto &item : m_whiteList)
    {
        ret = false;
        if (item.IPTbl.count(ipaddr & item.mask) > 0)
        {
            ret = true;
            break;
        }
    }
    if (m_workerType == WORKER_TYPE_PROXY)
    {
        if (ret && !m_identification.empty() && !username.empty() && username.length() > 0 && !password.empty() && password.length() > 0)
        {
            LOG_DEBUG("DoAuth: username %s password %s", username.c_str(), password.c_str());
            ret = false;
            if (m_identification[username] == password)
            {
                ret = true;
            }
        }
    }
    else if (m_workerType == WORKER_TYPE_SPIDER)
    {
        // 可以根据需要定制校验算法
        if (!username.empty())
        {
            ret = true;
        }
        else
        {
            ret = false;
        }
    }
    else
    {
        ret = false;
    }
    return ret;
}

bool Worker::Start()
{
    if (m_listenPortStart == 0 || m_listenNum <= 0)
    {
        return true;
    }
    m_workerLoop = uv_loop_new();
    if (!m_workerLoop)
    {
        LOG_DEBUG("Worker: %d Loop create failed", m_index);
        exit(1);
    }
    sockaddr_in addr;
    m_sessionServer = (uv_tcp_t *)calloc(m_listenNum, sizeof(uv_tcp_t));
    for (int i = 0; i < m_listenNum; i++)
    {
        uv_tcp_init(m_workerLoop, &m_sessionServer[i]);
        uv_ip4_addr("0.0.0.0", m_listenPortStart + i, &addr);
        uv_tcp_bind(&m_sessionServer[i], (const struct sockaddr *)&addr, 0);
        uv_handle_set_data((uv_handle_t *)&m_sessionServer[i], this);
        uv_listen((uv_stream_t *)&m_sessionServer[i], g_TcpListenSize, _RecvSessionConnection);
    }
    if (m_spiderPort)
    {
        uv_tcp_init(m_workerLoop, &m_spiderServer);
        uv_ip4_addr("0.0.0.0", m_spiderPort, &addr);
        uv_tcp_bind(&m_spiderServer, (const struct sockaddr *)&addr, 0);
        uv_handle_set_data((uv_handle_t *)&m_spiderServer, this);
        uv_listen((uv_stream_t *)&m_spiderServer, g_TcpListenSize, _RecvSpiderConnection);
    }
    if (uv_thread_create(&m_workerTid, _Run_Loop, this) != 0)
    {
        LOG_ERROR("Worker create thread failed");
    }
    return true;
}

void Worker::Stop()
{
    for (auto &pair : m_sessionCache)
    {
        pair.second->Close(false);
    }
    uv_stop(m_workerLoop);
    if (m_workerTid > 0)
    {
        uv_thread_join(&m_workerTid);
    }
    if (m_sessionServer)
    {
        free(m_sessionServer);
    }
    m_sessionCache.clear();
}

void Worker::RunLoop()
{
    LOG_DEBUG("Worker: %u RunLoop", m_index);
    uv_run(m_workerLoop, UV_RUN_DEFAULT);
    return;
}

void Worker::OnRecvSessionConnection(uv_stream_t *pServer)
{
    LOG_DEBUG("OnRecvConnection");
    auto pSession = new Session();
    if (pSession == nullptr)
    {
        return;
    }
    if (!pSession->Init(this, pServer, m_sessionId, m_commandTimeout, m_idleTimeout))
    {
        LOG_ERROR("Session: %08x Init failed", m_sessionId);
        delete pSession;
        return;
    }
    if (++m_sessionId >= (m_index << 24 | 0xFFFFFF))
    {
        m_sessionId = m_index << 24 | 1;
    }
    return;
}

void Worker::OnRecvSpiderConnection(uv_stream_t *pServer)
{
    LOG_DEBUG("OnRecvSpiderConnection");
    auto backConn = new SpiderBackConn();
    backConn->m_targetTcp = (uv_tcp_t *)calloc(1, sizeof(uv_tcp_t));
    uv_tcp_init(m_workerLoop, backConn->m_targetTcp);
    if (uv_accept(pServer, (uv_stream_t *)backConn->m_targetTcp) == 0)
    {
        backConn->m_pWorker = this;
        uv_timer_init(m_workerLoop, &backConn->m_sessionTimer);
        uv_handle_set_data((uv_handle_t *)&backConn->m_sessionTimer, backConn);
        uv_handle_set_data((uv_handle_t *)backConn->m_targetTcp, backConn);
        uv_read_start((uv_stream_t *)backConn->m_targetTcp, _OnAllocBuff, _OnReadData);
        uv_timer_start(&backConn->m_sessionTimer, _OnSpiderBackConnTimeout, g_CommandTimeout, 0);
        return;
    }
    return;
}

bool Worker::AddSessionCache(Session *pNode)
{
    if (m_sessionCache.count(pNode->m_sessionId)) {
        return false;
    }
    m_sessionCache[pNode->m_sessionId] = pNode;
    return true;
}

void Worker::DelSessionCache(Session *pNode)
{
    m_sessionCache.erase(pNode->m_sessionId);
    return;
}

Session *Worker::FindSessionCache(uint32_t SessionId)
{
    if (m_sessionCache.count(SessionId)) {
        return m_sessionCache[SessionId];
    }
    return nullptr;
}

void Worker::ShowWorker() const
{
    std::cout << "================================================================" << std::endl;
    std::cout << "Worker Index: " << m_index << std::endl;
    std::cout << "Worker Type: " << m_workerType << std::endl;
    std::cout << "Listen Port Start: " << m_listenPortStart << std::endl;
    std::cout << "Listen Port End: " << m_listenPortEnd << std::endl;
    std::cout << "Spider Port: " << m_spiderPort << std::endl;
    std::cout << "Identification: " << std::endl;
    for (const auto &pair : m_identification)
    {
        std::cout << "username: " << pair.first << ", password: " << pair.second << std::endl;
    }
    std::cout << "White List: " << std::endl;
    for (const auto &maskIps : m_whiteList)
    {
        for (const auto &ip : maskIps.IPTbl)
        {
            LOG_INFO("Add WhiteList IP [%08x] MASK[%08x]", ip, maskIps.mask);
        }
    }
    std::cout << "================================================================" << std::endl;
}

void SpiderBackConn::DoConn()
{
    LOG_DEBUG("SpiderBackConn DoConn");
    uv_read_stop((uv_stream_t *)m_targetTcp);
    uint32_t SessionId = *(uint32_t *)m_cmdRecvBuf;
    auto pSession = m_pWorker->FindSessionCache(SessionId);
    if (pSession) {
        pSession->AttachTarget(m_targetTcp);        
        m_targetTcp = nullptr;
    }
    Close();
}

void SpiderBackConn::Close()
{
    if (m_targetTcp) {
        LOG_DEBUG("SpiderBackConn Close");
        uv_close((uv_handle_t *)m_targetTcp, nullptr);
    }
    uv_timer_stop(&m_sessionTimer);
    uv_close((uv_handle_t *)&m_sessionTimer, _SpiderBackConnClose);
}